import os
import pandas as pd
import numpy as np
from tqdm import tqdm
import cv2
import glob
import pydicom
from joblib import Parallel, delayed
from PIL import Image
import pandas.api.types

import torch
import torchvision.models as models
import torchvision.transforms as transforms
import torchvision.datasets as datasets
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.autograd import Variable
from torch.utils.data.dataset import Dataset
import torchvision
import torchvision.transforms as transforms


class CNNModel(nn.Module):
    def __init__(self):
        super().__init__()

        self.input = nn.Conv2d(1, 3, kernel_size=3)
        model = models.efficientnet_b0(weights='IMAGENET1K_V1')

        self.features = model.features
        self.avgpool = model.avgpool

        self.output = nn.Linear(1280, 13)  # 使用清洗的数据

    def forward(self, x):

        # extract features
        x = self.input(x)
        x = self.features(x)
        x = self.avgpool(x)
        x = torch.flatten(x, 1)
        x = self.output(x)
        return x


def dicom_to_image(dicom_image):
    pixel_array = dicom_image.pixel_array

    if dicom_image.PixelRepresentation == 1:
        bit_shift = dicom_image.BitsAllocated - dicom_image.BitsStored
        dtype = pixel_array.dtype
        new_array = (pixel_array << bit_shift).astype(dtype) >> bit_shift
        pixel_array = pydicom.pixel_data_handlers.util.apply_modality_lut(
            new_array, dicom_image)

    if dicom_image.PhotometricInterpretation == "MONOCHROME1":
        pixel_array = 1 - pixel_array

    # transform to hounsfield units
    intercept = dicom_image.RescaleIntercept
    slope = dicom_image.RescaleSlope
    pixel_array = pixel_array * slope + intercept

    # windowing
    window_center = int(dicom_image.WindowCenter)
    window_width = int(dicom_image.WindowWidth)
    img_min = window_center - window_width // 2
    img_max = window_center + window_width // 2
    pixel_array = pixel_array.copy()
    pixel_array[pixel_array < img_min] = img_min
    pixel_array[pixel_array > img_max] = img_max

    # normalization
    pixel_array = (pixel_array - pixel_array.min()) / \
        (pixel_array.max() - pixel_array.min())

    return (pixel_array * 255).astype(np.uint8)


test_pids_path = r'dataset/sample_submission.csv'
test_path = r'/home/fd_chen/kaggle/rasac/dataset/test_images/'

test_pids = pd.read_csv(test_pids_path)['patient_id'].values

# Define the transform
transform = transforms.Compose([
    transforms.Resize((512, 512)),  # Resize the image to 224x224 pixels
    transforms.RandomVerticalFlip(0.5),  # Resize the image to 224x224 pixels
    transforms.RandomHorizontalFlip(0.5),  # Resize the image to 224x224 pixels
    transforms.ToTensor()  # Convert the image to a PyTorch tensor
])


model = CNNModel()
model.load_state_dict(torch.load(
    "checkpoints/0.002_split_model_205.pth"), strict=False)
model.cuda()



RESIZE = True
SIZE = 256
TICK = 5

pred_list = []
for pid in test_pids[:]:    
    print(pid)
    pid_paths_dcm_paths = glob.glob(
        f'/kaggle/input/rsna-2023-abdominal-trauma-detection/test_images/{str(pid)}/*/*'
    )
    pid_paths_dcm_paths.sort()
    pid_paths_dcm_paths = np.random.choice(pid_paths_dcm_paths, 5)
#     pid_paths_dcm_paths = pid_paths_dcm_paths[-25:]
    imgs = [Image.fromarray(dicom_to_image(pydicom.read_file(x))) for x in pid_paths_dcm_paths]
    imgs = [transform(x) for x in imgs]

    imgs = torch.cat(imgs, 0)
    imgs = imgs[:, None, :, :]

    with torch.no_grad():
        imgs = imgs.cuda()
        output = model(imgs)[:, :-1]
        pred = torch.sigmoid(output).data.cpu().numpy().round(3)
        pred = pred.mean(0)

    pred_list.append(pred)

pred_df = pd.DataFrame(np.vstack(pred_list))
pred_df.columns = ['bowel_healthy', 'bowel_injury', 'extravasation_healthy',
       'extravasation_injury', 'kidney_healthy', 'kidney_low', 'kidney_high',
       'liver_healthy', 'liver_low', 'liver_high', 'spleen_healthy',
       'spleen_low', 'spleen_high']

pred_df['patient_id'] = test_pids[:]


pred_df[['patient_id', 'bowel_healthy', 'bowel_injury', 'extravasation_healthy',
       'extravasation_injury', 'kidney_healthy', 'kidney_low', 'kidney_high',
       'liver_healthy', 'liver_low', 'liver_high', 'spleen_healthy',
       'spleen_low', 'spleen_high']].to_csv('submission.csv', index=None)




# !head /kaggle/working/submission.csv
